import math
import random
from typing import Union

from .agent import Agent, FastRandomAgent
from .board import GameState, Move
from .gotypes import Player
from .utils import CoordsFromPoint


def move_format(x: Union[Move, Player]):
    if x is Player.black:
        return "B"
    if x is Player.white:
        return "W"
    if x.is_pass:
        return "pass"
    if x.is_resign:
        return "resign"
    return CoordsFromPoint(x.point)


def show_tree(node, indent="", max_depth=3):
    if max_depth < 0:
        return
    if node is None:
        return
    if node.parent is None:
        print(f"{indent}root")
    else:
        player = node.parent.game_state.next_player
        move = node.move
        print(
            f"{indent}{move_format(player)} {move_format(move)} {node.num_rollouts} {node.winning_frac(player)}"
        )
    for child in sorted(node.children, key=lambda n: n.num_rollouts, reverse=True):
        show_tree(child, indent + "  ", max_depth - 1)


class MCTSNode:
    def __init__(self, game_state: GameState, parent=None, move: Move = None) -> None:
        self.game_state = game_state
        self.parent = parent
        self.move = move
        self.win_counts = {Player.black: 0, Player.white: 0}
        self.num_rollouts = 0
        self.children: list["MCTSNode"] = []
        self.unvisited_moves = game_state.LegalMoves()

    def AddRandomChild(self):
        index = random.randint(0, len(self.unvisited_moves) - 1)
        new_move = self.unvisited_moves.pop(index)
        new_game_state = self.game_state.ApplyMove(new_move)
        new_node = MCTSNode(new_game_state, self, new_move)
        self.children.append(new_node)
        return new_node

    def RecordWin(self, winner):
        self.win_counts[winner] += 1
        self.num_rollouts += 1

    def CanAddChild(self):
        return len(self.unvisited_moves) > 0

    def IsTerminal(self):
        return self.game_state.IsOver()

    def WinnerFrac(self, player: Player):
        return float(self.win_counts[player]) / float(self.num_rollouts)


class MCTSAgent(Agent):
    def __init__(self, rounds, temperature) -> None:
        super().__init__()
        self.rounds = rounds
        self.temperature = temperature

    def SelectMove(self, game_state: GameState):
        root = MCTSNode(game_state)
        for i in range(self.rounds):
            node = root
            while not node.CanAddChild() and not node.IsTerminal():
                node = self.SelectChild(node)
            if node.CanAddChild():
                node = node.AddRandomChild()
            winner = self.SimulateRandomGame(node.game_state)
            while node is not None:
                node.RecordWin(winner)
                node = node.parent
        scored_moves = [
            (child.WinnerFrac(game_state.next_player), child.move, child.num_rollouts)
            for child in root.children
        ]
        scored_moves.sort(key=lambda x: x[0], reverse=True)
        for s, m, n in scored_moves[:10]:
            print(f'{m} - {s:.3f} ({n})')
        best_move = None 
        best_pct = -1.0
        for child in root.children:
            child_pct = child.WinnerFrac(game_state.next_player)
            if child_pct > best_pct:
                best_pct = child_pct
                best_move = child.move
        print(f'Select move {best_move} with win pct {best_pct:.3f}')
        return best_move
    
    def SelectChild(self, node:MCTSNode):
        total_rollouts = sum(child.num_rollouts for child in node.children)
        log_rollouts = math.log(total_rollouts)
        best_score = -1
        best_child = None 
        for child in node.children:
            win_percentage = child.WinnerFrac(node.game_state.next_player)
            exploration_factor = math.sqrt(log_rollouts /child.num_rollouts)
            uct_score = win_percentage + self.temperature * exploration_factor
            if uct_score > best_score:
                best_score = uct_score
                best_child = child 
        return best_child
    
    @staticmethod
    def SimulateRandomGame(game_state:GameState):
        bots = {Player.black: FastRandomAgent(), Player.white: FastRandomAgent()}
        while not game_state.IsOver():
            move = bots[game_state.next_player].SelectMove(game_state)
            game_state = game_state.ApplyMove(move)
        return game_state.Winner()

